from rknn.api import RKNN
import onnx
import os

def convert_yolov8_to_rknn():
    model_path = './best.onnx'
    
    if not os.path.exists(model_path):
        print("模型文件不存在")
        return False
    
    # 创建RKNN对象
    rknn = RKNN(verbose=True)
    
    try:
        # 1. 基本配置（不使用量化相关参数）
        print('--> Configuring model')
        ret = rknn.config(
            mean_values=[[0, 0, 0]],
            std_values=[[255, 255, 255]],
            target_platform='rk3588'
            # 移除所有量化相关参数
        )
        if ret != 0:
            raise Exception(f'Config failed: {ret}')
        
        # 2. 加载ONNX
        print('--> Loading ONNX model')
        ret = rknn.load_onnx(model=model_path)
        if ret != 0:
            raise Exception(f'Load ONNX failed: {ret}')
        
        # 3. 构建模型（关闭量化）
        print('--> Building model')
        ret = rknn.build(do_quantization=False)  # 确保这里是False
        if ret != 0:
            raise Exception(f'Build failed: {ret}')
        
        # 4. 导出模型
        print('--> Exporting RKNN model')
        ret = rknn.export_rknn('yolov8s.rknn')
        if ret != 0:
            raise Exception(f'Export failed: {ret}')
        
        print('✓ RKNN模型转换成功!')
        return True
        
    except Exception as e:
        print(f'✗ 错误: {e}')
        return False
    finally:
        rknn.release()

if __name__ == '__main__':
    success = convert_yolov8_to_rknn()
    print("最终结果:", "成功" if success else "失败")
